{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a02706cb",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-10-21T05:52:06.240629Z",
     "start_time": "2022-10-21T05:52:03.394211Z"
    }
   },
   "source": [
    "import os\n",
    "import random\n",
    "import torch\n",
    "import dltools"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a16faf51",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-10-21T06:10:40.245918Z",
     "start_time": "2022-10-21T06:10:40.225972Z"
    }
   },
   "source": [
    "def _read_wiki(data_dir):\n",
    "    file_name = os.path.join(data_dir, 'wiki.train.tokens')\n",
    "    with open(file_name, 'r',encoding=\"utf-8\") as f:\n",
    "        lines = f.readlines()\n",
    "    # 大写字母转换为小写字母\n",
    "    paragraphs = [line.strip().lower().split(' . ') for line in lines if len(line.split(' . ')) >= 2]\n",
    "    random.shuffle(paragraphs)\n",
    "    return paragraphs"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "0c4da50f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-10-21T06:10:58.178352Z",
     "start_time": "2022-10-21T06:10:57.930511Z"
    }
   },
   "source": [
    "_read_wiki('./wikitext-2')"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f13a408d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-10-21T06:13:26.717452Z",
     "start_time": "2022-10-21T06:13:26.698503Z"
    }
   },
   "source": [
    "# 生成下一句预测任务的数据\n",
    "def _get_next_sentence(sentence, next_sentence, paragraphs):\n",
    "    if random.random() < 0.5:\n",
    "        is_next = True\n",
    "    else:\n",
    "        # paragraphs是三重列表的嵌套\n",
    "        next_sentence = random.choice(random.choice(paragraphs))\n",
    "        is_next = False\n",
    "    return sentence, next_sentence, is_next"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4db70d6d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-10-21T05:52:06.288501Z",
     "start_time": "2022-10-21T05:52:06.274536Z"
    }
   },
   "source": [
    "#@save\n",
    "def _get_nsp_data_from_paragraph(paragraph, paragraphs, vocab, max_len):\n",
    "    nsp_data_from_paragraph = []\n",
    "    for i in range(len(paragraph) - 1):\n",
    "        tokens_a, tokens_b, is_next = _get_next_sentence(\n",
    "            paragraph[i], paragraph[i + 1], paragraphs)\n",
    "        # 考虑1个'<cls>'词元和2个'<sep>'词元\n",
    "        if len(tokens_a) + len(tokens_b) + 3 > max_len:\n",
    "            continue\n",
    "        tokens, segments = dltools.get_tokens_and_segments(tokens_a, tokens_b)\n",
    "        nsp_data_from_paragraph.append((tokens, segments, is_next))\n",
    "    return nsp_data_from_paragraph"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "040bc80d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-10-21T06:21:48.602274Z",
     "start_time": "2022-10-21T06:21:48.594296Z"
    }
   },
   "source": [
    "# 生成遮蔽语言模型任务的数据\n",
    "def _replace_mlm_tokens(tokens, candidate_pred_positions, num_mlm_preds,\n",
    "                        vocab):\n",
    "    # 为遮蔽语言模型的输入创建新的词元副本，其中输入可能包含替换的“<mask>”或随机词元\n",
    "    mlm_input_tokens = [token for token in tokens]\n",
    "    pred_positions_and_labels = []\n",
    "    # 打乱后用于在遮蔽语言模型任务中获取15%的随机词元进行预测\n",
    "    random.shuffle(candidate_pred_positions)\n",
    "    for mlm_pred_position in candidate_pred_positions:\n",
    "        if len(pred_positions_and_labels) >= num_mlm_preds:\n",
    "            break\n",
    "        masked_token = None\n",
    "        # 80%的时间：将词替换为“<mask>”词元\n",
    "        if random.random() < 0.8:\n",
    "            masked_token = '<mask>'\n",
    "        else:\n",
    "            # 10%的时间：保持词不变\n",
    "            if random.random() < 0.5:\n",
    "                masked_token = tokens[mlm_pred_position]\n",
    "            # 10%的时间：用随机词替换该词\n",
    "            else:\n",
    "                masked_token = random.choice(vocab.idx_to_token)\n",
    "        mlm_input_tokens[mlm_pred_position] = masked_token\n",
    "        pred_positions_and_labels.append(\n",
    "            (mlm_pred_position, tokens[mlm_pred_position]))\n",
    "    return mlm_input_tokens, pred_positions_and_labels"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "de98eb80",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-10-21T05:52:06.319416Z",
     "start_time": "2022-10-21T05:52:06.304457Z"
    }
   },
   "source": [
    "# \n",
    "def _get_mlm_data_from_tokens(tokens, vocab):\n",
    "    candidate_pred_positions = []\n",
    "    # tokens是一个字符串列表\n",
    "    for i, token in enumerate(tokens):\n",
    "        # 在遮蔽语言模型任务中不会预测特殊词元\n",
    "        if token in ['<cls>', '<sep>']:\n",
    "            continue\n",
    "        candidate_pred_positions.append(i)\n",
    "    # 遮蔽语言模型任务中预测15%的随机词元\n",
    "    num_mlm_preds = max(1, round(len(tokens) * 0.15))\n",
    "    mlm_input_tokens, pred_positions_and_labels = _replace_mlm_tokens(\n",
    "        tokens, candidate_pred_positions, num_mlm_preds, vocab)\n",
    "    pred_positions_and_labels = sorted(pred_positions_and_labels,\n",
    "                                       key=lambda x: x[0])\n",
    "    pred_positions = [v[0] for v in pred_positions_and_labels]\n",
    "    mlm_pred_labels = [v[1] for v in pred_positions_and_labels]\n",
    "    return vocab[mlm_input_tokens], pred_positions, vocab[mlm_pred_labels]"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "6cb40b9d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-10-21T06:30:37.209988Z",
     "start_time": "2022-10-21T06:30:37.195027Z"
    }
   },
   "source": [
    "# 将文本转换为预训练数据集\n",
    "def _pad_bert_inputs(examples, max_len, vocab):\n",
    "    max_num_mlm_preds = round(max_len * 0.15)\n",
    "    all_token_ids, all_segments, valid_lens,  = [], [], []\n",
    "    all_pred_positions, all_mlm_weights, all_mlm_labels = [], [], []\n",
    "    nsp_labels = []\n",
    "    for (token_ids, pred_positions, mlm_pred_label_ids, segments,\n",
    "         is_next) in examples:\n",
    "        all_token_ids.append(torch.tensor(token_ids + [vocab['<pad>']] * (\n",
    "            max_len - len(token_ids)), dtype=torch.long))\n",
    "        all_segments.append(torch.tensor(segments + [0] * (\n",
    "            max_len - len(segments)), dtype=torch.long))\n",
    "        # valid_lens不包括'<pad>'的计数\n",
    "        valid_lens.append(torch.tensor(len(token_ids), dtype=torch.float32))\n",
    "        all_pred_positions.append(torch.tensor(pred_positions + [0] * (\n",
    "            max_num_mlm_preds - len(pred_positions)), dtype=torch.long))\n",
    "        # 填充词元的预测将通过乘以0权重在损失中过滤掉\n",
    "        all_mlm_weights.append(\n",
    "            torch.tensor([1.0] * len(mlm_pred_label_ids) + [0.0] * (\n",
    "                max_num_mlm_preds - len(pred_positions)),\n",
    "                dtype=torch.float32))\n",
    "        all_mlm_labels.append(torch.tensor(mlm_pred_label_ids + [0] * (\n",
    "            max_num_mlm_preds - len(mlm_pred_label_ids)), dtype=torch.long))\n",
    "        nsp_labels.append(torch.tensor(is_next, dtype=torch.long))\n",
    "    return (all_token_ids, all_segments, valid_lens, all_pred_positions,\n",
    "            all_mlm_weights, all_mlm_labels, nsp_labels)"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3f4738e0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-10-21T05:52:06.351333Z",
     "start_time": "2022-10-21T05:52:06.336371Z"
    }
   },
   "source": [
    "#@save\n",
    "class _WikiTextDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, paragraphs, max_len):\n",
    "        # 输入paragraphs[i]是代表段落的句子字符串列表；\n",
    "        # 而输出paragraphs[i]是代表段落的句子列表，其中每个句子都是词元列表\n",
    "        paragraphs = [dltools.tokenize(\n",
    "            paragraph, token='word') for paragraph in paragraphs]\n",
    "        sentences = [sentence for paragraph in paragraphs\n",
    "                     for sentence in paragraph]\n",
    "        self.vocab = dltools.Vocab(sentences, min_freq=5, reserved_tokens=[\n",
    "            '<pad>', '<mask>', '<cls>', '<sep>'])\n",
    "        # 获取下一句子预测任务的数据\n",
    "        examples = []\n",
    "        for paragraph in paragraphs:\n",
    "            examples.extend(_get_nsp_data_from_paragraph(\n",
    "                paragraph, paragraphs, self.vocab, max_len))\n",
    "        # 获取遮蔽语言模型任务的数据\n",
    "        examples = [(_get_mlm_data_from_tokens(tokens, self.vocab)\n",
    "                      + (segments, is_next))\n",
    "                     for tokens, segments, is_next in examples]\n",
    "        # 填充输入\n",
    "        (self.all_token_ids, self.all_segments, self.valid_lens,\n",
    "         self.all_pred_positions, self.all_mlm_weights,\n",
    "         self.all_mlm_labels, self.nsp_labels) = _pad_bert_inputs(\n",
    "            examples, max_len, self.vocab)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return (self.all_token_ids[idx], self.all_segments[idx],\n",
    "                self.valid_lens[idx], self.all_pred_positions[idx],\n",
    "                self.all_mlm_weights[idx], self.all_mlm_labels[idx],\n",
    "                self.nsp_labels[idx])\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.all_token_ids)"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1cf21293",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-10-21T05:52:06.366292Z",
     "start_time": "2022-10-21T05:52:06.352328Z"
    }
   },
   "source": [
    "#@save\n",
    "def load_data_wiki(batch_size, max_len):\n",
    "    \"\"\"加载WikiText-2数据集\"\"\"\n",
    "    num_workers = dltools.get_dataloader_workers()\n",
    "    data_dir = \"./wikitext-2/\"\n",
    "    paragraphs = _read_wiki(data_dir)\n",
    "    train_set = _WikiTextDataset(paragraphs, max_len)\n",
    "    train_iter = torch.utils.data.DataLoader(train_set, batch_size,\n",
    "                                        shuffle=True, num_workers=num_workers)\n",
    "    return train_iter, train_set.vocab"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e041e570",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-10-21T05:52:12.097637Z",
     "start_time": "2022-10-21T05:52:06.369283Z"
    }
   },
   "source": [
    "batch_size, max_len = 512, 64\n",
    "train_iter, vocab = load_data_wiki(batch_size, max_len)\n",
    "\n",
    "for (tokens_X, segments_X, valid_lens_x, pred_positions_X, mlm_weights_X,\n",
    "     mlm_Y, nsp_y) in train_iter:\n",
    "    print(tokens_X.shape, segments_X.shape, valid_lens_x.shape,\n",
    "          pred_positions_X.shape, mlm_weights_X.shape, mlm_Y.shape,\n",
    "          nsp_y.shape)\n",
    "    break"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "344a00ae",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-10-21T05:52:12.113594Z",
     "start_time": "2022-10-21T05:52:12.099632Z"
    }
   },
   "source": [
    "len(vocab)"
   ],
   "outputs": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.8"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
